外部ネットワークアクセスで RDS PostgreSQL に IAM認証で接続してみた #SnowflakeDB

-- データベースを作成
CREATE DATABASE sample_db;

\c sample_db;

-- テーブル: customers 
CREATE TABLE customers (
    customer_id SERIAL PRIMARY KEY,  
    name VARCHAR(100),               
    email VARCHAR(100),              
    age INT                          
);

-- customers テーブルにサンプルデータを挿入
INSERT INTO customers (name, email, age) VALUES
('Alice Smith', 'alice@example.com', 30),
('Bob Johnson', 'bob@example.com', 25),
('Charlie Brown', 'charlie@example.com', 35);

--専用のロールを作成
CREATE ROLE readonly_role;
GRANT CONNECT ON DATABASE sample_db TO readonly_role;

\c sample_db;
GRANT USAGE ON SCHEMA public TO readonly_role;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO readonly_role;

-- IAMデータベース認証用のユーザーを作成し権限を付与
CREATE USER iam_ro_user WITH LOGIN;
GRANT rds_iam TO iam_ro_user;
GRANT readonly_role TO iam_ro_user;

{
   "Version": "2012-10-17",
   "Statement": [
      {
         "Effect": "Allow",
         "Action": [
             "rds-db:connect"
         ],
         "Resource": [
             "arn:aws:rds-db:ap-northeast-1:<account-id>:dbuser:<dbi-resource-id>/iam_ro_user"
         ]
      }
   ]
}

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": "sts:AssumeRole",
            "Principal": {
                "AWS": "<自身のAWSアカウントID>"
            },
            "Condition": {
                "StringEquals": {
                    "sts:ExternalId": "<任意の値>"
                }
            }
        }
    ]
}

CREATE OR REPLACE NETWORK RULE aws_rds_postgres_network_rule
    MODE = EGRESS
    TYPE = PRIVATE_HOST_PORT
    VALUE_LIST = ('database-1.xxxxx.ap-northeast-1.rds.amazonaws.com:5432');

CREATE OR REPLACE SECURITY INTEGRATION rds_security_integration
  TYPE = API_AUTHENTICATION
  AUTH_TYPE = AWS_IAM
  ENABLED = TRUE
  AWS_ROLE_ARN = 'arn:aws:iam::<account-id>:role/sf-udf-rds-auth-role';

DESC  SECURITY INTEGRATION rds_security_integration;

{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Effect": "Allow",
			"Principal": {
				"AWS": "<API_AWS_IAM_USER_ARN>"
			},
			"Action": "sts:AssumeRole",
			"Condition": {
				"StringEquals": {
					"sts:ExternalId": "<API_AWS_EXTERNAL_ID>"
				}
			}
		}
	]
}

CREATE OR REPLACE SECRET aws_rds_access_token
  TYPE = CLOUD_PROVIDER_TOKEN
  API_AUTHENTICATION = rds_security_integration;

--外部ネットワークアクセスを作成
CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION rds_external_access_integration
  ALLOWED_NETWORK_RULES = (<db>.<schema>.aws_rds_postgres_network_rule)
  ALLOWED_AUTHENTICATION_SECRETS=(<db>.<schema>.aws_rds_access_token)
  ENABLED=true ;

CREATE OR REPLACE FUNCTION fetch_temp_credentials()
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
HANDLER = 'get_secret_username_password'
EXTERNAL_ACCESS_INTEGRATIONS = (rds_iam_external_access_integration)
SECRETS = ('cred' = <db>.<schema>.aws_rds_access_token)
AS
$$
import _snowflake
import json

def get_secret_username_password():
  cloud_provider_object = _snowflake.get_cloud_provider_token('cred')
  cloud_provider_dictionary = {
          "ACCESS_KEY_ID": cloud_provider_object.access_key_id,
          "SECRET_ACCESS_KEY": cloud_provider_object.secret_access_key,
          "TOKEN": cloud_provider_object.token
      }
  return json.dumps(cloud_provider_dictionary)
$$;

{
	"ACCESS_KEY_ID": "xxxxxxxxxx", 
	"SECRET_ACCESS_KEY": "xxxxx", 
	"TOKEN": "xxxxxxxxxxxxxxxxxx"
}

CREATE OR REPLACE function query_postgres()
RETURNS VARIANT
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
IMPORTS=('@my_int_stage/ap-northeast-1-bundle.pem')
PACKAGES = ('psycopg2','boto3')
HANDLER = 'main'
EXTERNAL_ACCESS_INTEGRATIONS = (rds_iam_external_access_integration)
SECRETS = ('cred' = <db>.<schema>.aws_rds_access_token)
AS $$
import psycopg2
import sys
import os
import json
import boto3
import _snowflake

# RDS PostgreSQL の情報を設定
ENDPOINT = "database-1.xxxxx.ap-northeast-1.rds.amazonaws.com"
DBNAME = "sample_db"
USER = "iam_ro_user"
PORT = 5432
REGION = "ap-northeast-1"

# Snowflake UDFのインポートディレクトリから証明書ファイルのパスを取得
IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
ssl_cert_path = os.path.join(import_dir, 'ap-northeast-1-bundle.pem')  # 証明書ファイルのパス

def main():
    # シークレットから一時的な認証情報を取得
    cloud_provider_object = _snowflake.get_cloud_provider_token('cred')
    cloud_provider_dictionary = {
          "ACCESS_KEY_ID": cloud_provider_object.access_key_id,
          "SECRET_ACCESS_KEY": cloud_provider_object.secret_access_key,
          "TOKEN": cloud_provider_object.token
      }

    # 認証情報の割り当てとリージョンを指定
    boto3_session_args = {
          'aws_access_key_id': cloud_provider_dictionary["ACCESS_KEY_ID"],
          'aws_secret_access_key': cloud_provider_dictionary["SECRET_ACCESS_KEY"],
          'aws_session_token': cloud_provider_dictionary["TOKEN"],
          'region_name': 'ap-northeast-1'
      }

    # セッションを作成
    session = boto3.Session(**boto3_session_args)

    # RDS を操作するためのクライアントを作成
    client = session.client('rds')

    # RDS が発行する認証トークンを取得
    token = client.generate_db_auth_token(DBHostname=ENDPOINT, Port=PORT, DBUsername=USER, Region=REGION)

    # PostgreSQLに接続
    conn = psycopg2.connect(
        host=ENDPOINT,
        port=PORT,
        database=DBNAME,
        user=USER,
        password=token,
        sslmode='verify-ca',
        sslrootcert=ssl_cert_path # インポートディレクトリから取得した証明書ファイルを指定
    )

    # クエリを実行
    cur = conn.cursor()
    cur.execute("SELECT * FROM customers")
    result = cur.fetchall()

    # 接続を閉じる
    conn.close()
    return result
$$;